#! /usr/bin/env python
#
# Python implementation of the demo functions provided in cookbook.c that
# comes with CFITSIO
#
# TODO:
#   - check doc strings
#
import sys
import os

import ctypes

# Use just the C-types bindings to CFITSIO
from fitsio import cfitsio

def printerror(status):
    '''Print out CFITSIO error message and exit program

    @param status: FITSIO status.
    @type status: L{ctypes.c_int}
    '''

    # Not done exactly the way as in cookbook.c

    errtext = ctypes.create_string_buffer('\000' * cfitsio.FLEN_STATUS)
    cfitsio.ffgerr(status, errtext)
    print status.value, errtext.value

    errmsg = ctypes.create_string_buffer('\000' * cfitsio.FLEN_ERRMSG)
    while cfitsio.fits_read_errmsg(errmsg):
        print errmsg.value

    # Note: the pyfitsio module provides a more pythonic error handler that
    # uses exceptions
    sys.exit(status.value)

    return

def remove(filename):
    '''Delete file if it already exists

    @param filename: Name of file to delete
    @type filename: python string
    '''
    if os.path.exists(filename):
        os.remove(filename)

def make_c_array(ctype, *args):
    '''Convenience function to create and initialize a ctypes array

    @param ctype: C-types primitive specifying array type
    @type ctype: instance of ctypes type

    @param args: list of initial values for the array
    @type args: python tuple
    '''
    size = len(args)
    return (ctype * size)(*args)

#------------------------------------------------------------------------------

def writeimage():
    '''
    Create a FITS primary array containing a 2-D image
    '''

    filename = 'atestfil.fit'

    # Delete file if it already exists
    remove(filename)

    # A two dimensional image
    naxis = ctypes.c_int(2)

    # Image is 300 by 200 pixels
    naxes = (ctypes.c_long * 2)(300,200)

    # Allocate memory for the whole image
    array = ( ctypes.c_short * (naxes[0] * naxes[1]) )()

    # Get a pointer to the FITS file
    fptr = ctypes.pointer( cfitsio.fitsfile() )

    # Initialize status to zero
    status = ctypes.c_int(0)

    # Create new FITS file, note that we need a pointer-to-pointer to the
    # fitsfile
    if cfitsio.fits_create_file(ctypes.pointer(fptr), filename, \
                                ctypes.byref(status)):
        printerror(status)

    # Write the required keywords for the primary array image which will be
    # of unsigned short
    bitpix = cfitsio.USHORT_IMG
    if cfitsio.fits_create_img(fptr, bitpix, naxis, naxes, ctypes.byref(status)):
        printerror(status)

    # Fill up image with linear ramp function (not exactly as in cookbook)
    for j in range(naxes[1]):
        for i in range(naxes[0]):
            array[j * naxes[0] + i] = i + j

    # First pixel to write
    fpixel = ctypes.c_longlong(1)

    # Number of pixels to write
    nelements = ctypes.c_longlong(naxes[0] * naxes[1])

    # Write the array of unsigned integers to the FITS file. Note that
    # fixel and nelements need tobe LONGLONGs
    if cfitsio.fits_write_img(fptr, cfitsio.TUSHORT, fpixel, nelements, array,
                     ctypes.byref(status)):
        printerror(status)

    # Write another optional keyword to the header
    exposure = ctypes.c_long(1500)
    if cfitsio.fits_update_key(fptr, cfitsio.TLONG, "EXPOSURE", \
                               ctypes.byref(exposure),
                     "Total Exposure Time", ctypes.byref(status)):
        printerror(status)

    # Close the file
    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)
    
    return

#------------------------------------------------------------------------------

def writeascii():
    '''
    Create an ASCII table extension containing 3 columns and 6 rows
    '''

    filename = 'atestfil.fit'
    extname = 'PLANETS_ASCII'

    # ctypes equivalent of char*[]
    ttype = make_c_array(ctypes.c_char_p, 'Planet', 'Diameter', 'Density')
    tform = make_c_array(ctypes.c_char_p, 'a8', 'I6', 'F4.2')
    tunit = make_c_array(ctypes.c_char_p, '\0', 'km', 'g/cm')

    #ttype = (ctypes.c_char_p * 3)('Planet', 'Diameter', 'Density')

    # Define the values for each planet
    planet = make_c_array(ctypes.c_char_p, \
                          'Mercury', 'Venus', 'Earth', \
                          'Mars', 'Jupiter', 'Saturn')
    diameter = make_c_array(ctypes.c_long, \
                        4880, 12112, 12742, 6800, 143000, 121000)
    density  = make_c_array(ctypes.c_float, 5.1, 5.3, 5.52, 3.94, 1.33, 0.69)

    fptr = ctypes.pointer( cfitsio.fitsfile() )
    status = ctypes.c_int(0)

    if cfitsio.fits_open_file(ctypes.pointer(fptr), filename, \
                              cfitsio.READWRITE, ctypes.byref(status)):
        printerror(status)

    # Create empty ascii table with columns correctly labeled
    tfields = ctypes.c_int(3)
    nrows   = ctypes.c_longlong(6)
    if cfitsio.fits_create_tbl(fptr, cfitsio.ASCII_TBL, nrows, \
                               tfields, ttype, tform, tunit, extname, \
                               ctypes.byref(status)):
        printerror(status)

    firstrow  = ctypes.c_longlong(1)
    firstelem = ctypes.c_longlong(1)

    # Write names as strings to the first column
    if cfitsio.fits_write_col(fptr, cfitsio.TSTRING, 1, firstrow, firstelem, \
                              nrows, planet, ctypes.byref(status)):
        printerror(status)

    # Write diameters as longs to the second column
    if cfitsio.fits_write_col(fptr, cfitsio.TLONG, 2, firstrow, firstelem, \
                              nrows, diameter, ctypes.byref(status)):
        printerror(status)

    # Write density as floats to the third column
    if cfitsio.fits_write_col(fptr, cfitsio.TFLOAT, 3, firstrow, firstelem, \
                              nrows, density, ctypes.byref(status)):
        printerror(status)        

    # Close the file
    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)

#------------------------------------------------------------------------------

def writebintable():
    '''
    Create a binary table extension containing 3 columns and 6 rows
    '''

    filename = 'atestfil.fit'
    extname = 'PLANETS_Binary'
    
    # ctypes equivalent of char*[]
    ttype = make_c_array(ctypes.c_char_p, 'Planet', 'Diameter', 'Density')
    tform = make_c_array(ctypes.c_char_p, '8a', '1J', '1E')
    tunit = make_c_array(ctypes.c_char_p, '\0', 'km', 'g/cm')

    # Define the values for each planet
    planet = make_c_array(ctypes.c_char_p, \
                          'Mercury', 'Venus', 'Earth', \
                          'Mars', 'Jupiter', 'Saturn')
    diameter = make_c_array(ctypes.c_long, \
                            4880, 12112, 12742, 6800, 143000, 121000)
    density  = make_c_array(ctypes.c_float, 5.1, 5.3, 5.52, 3.94, 1.33, 0.69)
    
    
    fptr = ctypes.pointer( cfitsio.fitsfile() )
    status = ctypes.c_int(0)

    if cfitsio.fits_open_file(ctypes.pointer(fptr), filename, \
                              cfitsio.READWRITE, ctypes.byref(status)):
        printerror(status)

    # Move to 2nd HDU
    hdutype = ctypes.c_int()
    if cfitsio.fits_movabs_hdu(fptr, 2, ctypes.byref(hdutype), \
                               ctypes.byref(status)):
        printerror(status)

    # Append a new empty binary table
    tfields = ctypes.c_int(3)
    nrows   = ctypes.c_longlong(6)
    if cfitsio.fits_create_tbl(fptr, cfitsio.BINARY_TBL, nrows, \
                               tfields, ttype, tform, tunit, extname, \
                               ctypes.byref(status)):
        printerror(status)
        
    firstrow  = ctypes.c_longlong(1)
    firstelem = ctypes.c_longlong(1)

    # Write names as strings to the first column
    if cfitsio.fits_write_col(fptr, cfitsio.TSTRING, 1, firstrow, firstelem, \
                              nrows, planet, ctypes.byref(status)):
        printerror(status)

    # Write diameters as longs to the second column
    if cfitsio.fits_write_col(fptr, cfitsio.TLONG, 2, firstrow, firstelem, \
                              nrows, diameter, ctypes.byref(status)):
        printerror(status)

    # Write density as floats to the third column
    if cfitsio.fits_write_col(fptr, cfitsio.TFLOAT, 3, firstrow, firstelem, \
                              nrows, density, ctypes.byref(status)):
        printerror(status)        

    # Close the file
    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)

#------------------------------------------------------------------------------

def copyhdu():
    '''
    Copy 1st and 3rd HDUs from input file to a new FITS file
    '''

    infptr  = ctypes.pointer( cfitsio.fitsfile() )
    outfptr = ctypes.pointer( cfitsio.fitsfile() )

    infilename  = 'atestfil.fit'
    outfilename = 'btestfil.fit'

    # Delete old output file if it already exists
    remove(outfilename)

    status = ctypes.c_int(0)

    # Open the existing FITS file
    if cfitsio.fits_open_file(ctypes.pointer(infptr), infilename, \
                              cfitsio.READONLY, ctypes.byref(status)):
        printerror(status)

    # Create another FITS file
    if cfitsio.fits_create_file(ctypes.pointer(outfptr), outfilename, \
                                ctypes.byref(status)):
        printerror(status)

    # Copy primary array from the input to output file
    morekeys = ctypes.c_int(0)
    if cfitsio.fits_copy_hdu(infptr, outfptr, morekeys, ctypes.byref(status)):
        printerror(status)

    # Move to 3rd HDU in the input file
    hdutype = ctypes.c_int()
    if cfitsio.fits_movabs_hdu(infptr, 3, ctypes.byref(hdutype), \
                               ctypes.byref(status)):
        printerror(status)

    # Copy 3rd HDU from input file to the 2nd HDU of output file
    if cfitsio.fits_copy_hdu(infptr, outfptr, morekeys, ctypes.byref(status)):
        printerror(status)

    # Close both of the FITS files
    if cfitsio.fits_close_file(infptr, ctypes.byref(status)):
        printerror(status)

    if cfitsio.fits_close_file(outfptr, ctypes.byref(status)):
        printerror(status)                                                        
    return
#------------------------------------------------------------------------------

def selectrows():
    '''
    Select rows from an input table and copy to the output table
    '''

    infptr  = ctypes.pointer( cfitsio.fitsfile() )
    outfptr = ctypes.pointer( cfitsio.fitsfile() )

    infilename  = 'atestfil.fit'
    outfilename = 'btestfil.fit'

    status = ctypes.c_int(0)

    # Open the existing FITS files
    if cfitsio.fits_open_file(ctypes.pointer(infptr), infilename, \
                              cfitsio.READONLY, ctypes.byref(status)):
        printerror(status)
    if cfitsio.fits_open_file(ctypes.pointer(outfptr), outfilename, \
                              cfitsio.READWRITE, ctypes.byref(status)):
        printerror(status)

    # Move to 3rd HDU in the input file
    hdutype = ctypes.c_int()
    if cfitsio.fits_movabs_hdu(infptr, 3, ctypes.byref(hdutype), \
                               ctypes.byref(status)):
        printerror(status)

    if not hdutype.value == cfitsio.BINARY_TBL:
        print 'Error: expected to find a binary table in this HDU'
        sys.exit(1)

    # Move to the 2nd HDU in output file
    if cfitsio.fits_movabs_hdu(outfptr, 2, ctypes.byref(hdutype), \
                               ctypes.byref(status)):
        printerror(status)

    # Create new extension in the output file
    if cfitsio.fits_create_hdu(outfptr, ctypes.byref(status)):
        printerror(status)

    # Get number of keywords
    nkeys = ctypes.c_int()
    keypos = ctypes.c_int()
    if cfitsio.fits_get_hdrpos(infptr, ctypes.byref(nkeys), \
                               ctypes.byref(keypos), ctypes.byref(status)):
        printerror(status)

    # Copy all keywords from input to the output extension
    card = ctypes.create_string_buffer('\000' * cfitsio.FLEN_CARD)
    for i in range(1, nkeys.value+1):
        cfitsio.fits_read_record(infptr, i, card, ctypes.byref(status))
        cfitsio.fits_write_record(outfptr, card, ctypes.byref(status))
        #print i, card.value

    # Read NAXIS1 and NAXIS2 keywords to get table size
    naxes = (ctypes.c_long * 2)()
    nfound = ctypes.c_int()
    if cfitsio.fits_read_keys_lng(infptr, "NAXIS", 1, 2, naxes, \
                                  ctypes.byref(nfound),
                          ctypes.byref(status)):
        printerror(status)

    # Find which column cotains the density values
    colnum = ctypes.c_int()
    if cfitsio.fits_get_colnum(infptr, cfitsio.CASEINSEN, "density", \
                               ctypes.byref(colnum),
                       ctypes.byref(status)):
        printerror(status)

    # Read the density column values
    frow = ctypes.c_longlong(1)
    felem = ctypes.c_longlong(1)
    nullval = ctypes.c_float(-99.)
    density = (ctypes.c_float * naxes[1])()
    anynulls = ctypes.c_int()
    if cfitsio.fits_read_col(infptr, cfitsio.TFLOAT, colnum, frow, felem, \
                             naxes[1], ctypes.byref(nullval), density, \
                             ctypes.byref(anynulls), ctypes.byref(status)):
        printerror(status)

    # allocate buffer large enough for 1 row of the table
    buff = (ctypes.c_ubyte * naxes[0])()

    # If the density is less than 3.0, copy the row to the outpur table
    noutrows = ctypes.c_longlong(0)
    for irow in range(1, naxes[1]+1):
        if density[irow-1] < 3.0:
            noutrows.value += 1
            cfitsio.fits_read_tblbytes(infptr, irow, 1, naxes[0], buff, \
                                       ctypes.byref(status))
            cfitsio.fits_write_tblbytes(outfptr, noutrows, 1, naxes[0], buff, \
                                        ctypes.byref(status))
    
    # Update NAXIS2 keyword with the correct number of rows. Note use of
    # None to set comment argument as NULL pointer
    if cfitsio.fits_update_key(outfptr, cfitsio.TLONG, "NAXIS2", \
                               ctypes.byref(noutrows), None, \
                               ctypes.byref(status)):
        printerror(status)

    # Close both of the FITS files
    if cfitsio.ffclos(infptr, ctypes.byref(status)):
        printerror(status)
    if cfitsio.ffclos(outfptr, ctypes.byref(status)):
        printerror(status)

    return

#------------------------------------------------------------------------------

def readheader():
    '''
    Print out all header keywords in all extensions of a FITS file.
    '''

    filename = 'atestfil.fit'
    fptr = ctypes.pointer( cfitsio.fitsfile() )

    status = ctypes.c_int(0)

    if cfitsio.fits_open_file(ctypes.pointer(fptr), filename, \
                              cfitsio.READONLY, ctypes.byref(status)):
        printerror(status)

    # Create string for header record
    card = ctypes.create_string_buffer('\000' * cfitsio.FLEN_CARD)

    # Attempt to move to next HDU until we get an EOF error
    ii = 1
    hdutype = ctypes.c_int()
    while not cfitsio.fits_movabs_hdu(fptr, ii, ctypes.byref(hdutype), \
                                      ctypes.byref(status)):

        # Get number of keywords
        nkeys = ctypes.c_int()
        keypos = ctypes.c_int()
        if cfitsio.fits_get_hdrpos(fptr, ctypes.byref(nkeys), \
                                   ctypes.byref(keypos), ctypes.byref(status)):
            printerror(status)
        
        print "Header listing for HDU #", ii, nkeys, keypos
        for jj in range(1, nkeys.value+1):
            if cfitsio.fits_read_record(fptr, jj, card, ctypes.byref(status)):
                printerror(status)
            print card.value

        ii += 1

    # Check that we got the expected EOF and reset status
    if status.value == cfitsio.END_OF_FILE:
        status.value = 0
    else:
        printerror(status)
    
    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)
    
#------------------------------------------------------------------------------

def readimage():
    '''
    Read a FITS image and determine the minimum and maximum pixel values
    '''

    filename = 'atestfil.fit'
    fptr = ctypes.pointer( cfitsio.fitsfile() )

    status = ctypes.c_int(0)

    # Read the image 1000 bytes at a time
    buffsize = 1000
    buffer = (ctypes.c_float * buffsize)()

    if cfitsio.fits_open_file(ctypes.pointer(fptr), filename, \
                              cfitsio.READONLY, ctypes.byref(status)):
        printerror(status)

    # Read NAXIS1 and NAXIS2 keywords to get image size
    naxes = (ctypes.c_long * 2)()
    nfound = ctypes.c_int()
    if cfitsio.fits_read_keys_lng(fptr, "NAXIS", 1, 2, naxes, \
                                  ctypes.byref(nfound), ctypes.byref(status)):
        printerror(status)

    npixels = naxes[0] * naxes[1]
    fpixel  = ctypes.c_longlong(1)
    nullval = ctypes.c_float(0)
    datamin = 1.0e30
    datamax = -1.0e30
    anynull = ctypes.c_int()
    
    while (npixels > 0):

        # read as many pixels as will fit in the buffer
        nbuffer = ctypes.c_longlong(npixels)
        if npixels > buffsize:
            nbuffer = ctypes.c_longlong(buffsize)

        if cfitsio.fits_read_img(fptr, cfitsio.TFLOAT, fpixel, nbuffer, \
                                 ctypes.byref(nullval), buffer, \
                                 ctypes.byref(anynull), ctypes.byref(status)):
            printerror(status)
        
        for ii in range(nbuffer.value):
            if buffer[ii] < datamin:
                datamin = buffer[ii]
            if buffer[ii] > datamax:
                datamax = buffer[ii]

        npixels -= nbuffer.value
        fpixel.value += nbuffer.value

    print "\nMin and max image pixels =", datamin, datamax

    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)
        
#------------------------------------------------------------------------------

def readtable():
    '''
    Read and print data values from an ASCII or binary table
    '''
    
    filename = 'atestfil.fit'
    fptr = ctypes.pointer( cfitsio.fitsfile() )
    status = ctypes.c_int(0)

    if cfitsio.fits_open_file(ctypes.pointer(fptr), filename, \
                              cfitsio.READONLY, ctypes.byref(status)):
        printerror(status)

    # This is a bit of a FUDGE
    numcols = 3
    ttype = (ctypes.c_char_p * numcols)()
    for ii in range(numcols):
        temp = ctypes.create_string_buffer('\000' * cfitsio.FLEN_VALUE)
        temp.value = str(ii)
        ttype[ii] = temp.value

    # This is a FUDGE too
    name  = (ctypes.c_char_p * 10)()
    for ii in range(10):
        temp = ctypes.create_string_buffer('\000' * 10)
        temp.value = str(ii)
        name[ii] = temp.value

    # Read ASCII then binary table
    hdutype = ctypes.c_int()
    for hdunum in range(2, 4):

        if cfitsio.fits_movabs_hdu(fptr, hdunum, ctypes.byref(hdutype), \
                                   ctypes.byref(status)):
            printerror(status)

        if hdutype.value == cfitsio.ASCII_TBL:
            print "Reading ASCII table in HDU#", hdunum
        elif hdutype.value == cfitsio.BINARY_TBL:
            print "Reading binary table in HDU#", hdunum
        else:
            print "Error: this HDU is not an ASCII or binary table"
            sys.exit(-1)

        # Read the column names from the TTYPEn keywords
        nfound = ctypes.c_int()
        cfitsio.fits_read_keys_str(fptr, "TTYPE", 1, 3, ttype, \
                                   ctypes.byref(nfound),
                           ctypes.byref(status))
        print " Row %10s %10s %10s" % (ttype[0], ttype[1], ttype[2])


        frow = ctypes.c_longlong(1)
        felem = ctypes.c_longlong(1)
        nelem = ctypes.c_longlong(6)
        strnull = ' '
        longnull = ctypes.c_long(0)
        floatnull = ctypes.c_float(0.0)
        anynull = ctypes.c_int()
        
        cfitsio.fits_read_col(fptr, cfitsio.TSTRING, 1, frow, felem, nelem, \
                              strnull, name, ctypes.byref(anynull), \
                              ctypes.byref(status))

        dia = (ctypes.c_long * 6)()
        cfitsio.fits_read_col(fptr, cfitsio.TLONG, 2, frow, felem, nelem, \
                              strnull, dia, ctypes.byref(anynull), \
                              ctypes.byref(status))
        
        den = (ctypes.c_float * 6)()
        cfitsio.fits_read_col(fptr, cfitsio.TFLOAT, 3, frow, felem, nelem, \
                              strnull, den, ctypes.byref(anynull), \
                              ctypes.byref(status))
        
        for ii in range(6):
            print "%5d %10s %10ld %10.2f" % (ii+1, name[ii], dia[ii], den[ii])

    if cfitsio.fits_close_file(fptr, ctypes.byref(status)):
        printerror(status)
    
#------------------------------------------------------------------------------

def main():

    writeimage()
    writeascii()
    writebintable()
    copyhdu()
    selectrows()
    readheader()
    readimage()
    readtable()

    print "\nAll the pyfitsio routines were called successfully"

#------------------------------------------------------------------------------

if __name__ == '__main__':
    main()
